#!/usr/bin/env python3
'''
与atari-doubledunk train_a2c_lstm.py和atari-break train_a2c_lstm.py得代码进行对比
已适配

训练记录：
在2号机上训练
20250115:训练分数7.112，测试分数15.84分，继续训练
20250116：Learning Rate: 0.0005 scheduler last epoch:  13600，训练分数8.127分，测试分数19.8分，继续训练
20250117：Learning Rate: 0.0005，训练分数分数13.825分，测试分数19.8分，继续训练
20250119:Learning Rate: 0.0005，训练分数16.586，测试分数29.77分，继续训练
20250120：Learning Rate: 0.0005，因断电训练分数未知，测试分数29.77分，继续训练
20250121:Learning Rate: 0.00045000000000000004，训练分数25.927分，测试分数29.77分，继续训练
20250122: Learning Rate: 0.00045000000000000004，训练分数26.980，测试分数43.885分，继续训练
20250123：Learning Rate: 0.00045000000000000004，训练分数30.933，测试分数43.885分，继续训练
20250125:Learning Rate: 0.00045000000000000004，训练分数29.332，测试分数59.87，继续训练
20250126:Learning Rate: 0.00045000000000000004，训练分数92.602分，测试分数108.97分，继续训练，看来模型已经突破了前提探索，发现了正向策略的方向
20250127:Learning Rate: 0.00040500000000000003,scheduler last epoch:  121800，训练分数94.104，测试分数108.97分，继续训练
20250128:Learning Rate: 0.00040500000000000003,scheduler last epoch:  145000，训练分数98.921，测试分数108.99，继续训练
20250129:Learning Rate: 0.0003645，scheduler last epoch:  168600，训练分数93.543，测试分数108.99（新分数108.98）分，继续训练一天
20250130:Learning Rate: 0.0003645,epoch:172500,训练分数97.536，测试分数108.99分，停止训练，play模型
'''
import gymnasium as gym
import ptan
import numpy as np
import argparse
from tensorboardX import SummaryWriter
import os

import torch
import torch.nn as nn
import torch.nn.utils as nn_utils
import torch.nn.functional as F
import torch.optim as optim

from typing import Any
from lib import common, model
from collections import deque
import ale_py

gym.register_envs(ale_py)

GAMMA = 0.99
LEARNING_RATE = 5e-4
ENTROPY_BETA = 0.01
BATCH_SIZE = 128
NUM_ENVS = 50

REWARD_STEPS = 4
CLIP_GRAD = 0.5

SAVE_ITERS = 100


def unpack_batch(batch, net, device='cpu'):
    """
    Convert batch into training tensors
    :param batch:
    :param net:
    :return: states variable, actions tensor, reference values variable
    """
    states = []
    actions = []
    rewards = []
    not_done_idx = [] # 非结束的游戏数据索引，该索引记录对应batch，states，actions，rewards
    last_states = [] # 记录采样中的执行动作后的状态，仅记录游戏非结束状态下的索引
    last_states_hx = []
    last_states_cx = []
    hx = []
    cx = []
    for idx, exp in enumerate(batch):
        states.append(np.asarray(exp[0]))
        actions.append(int(exp[1]))
        rewards.append(exp[2])
        if not exp[3]:
            not_done_idx.append(idx)
            last_states.append(np.asarray(exp[4]))
            last_states_hx.append(exp[5][0])
            last_states_cx.append(exp[5][1])
        hx.append(exp[5][0])
        cx.append(exp[5][1])
    
    states_v = torch.FloatTensor(np.asarray(states)).to(device)
    actions_t = torch.LongTensor(actions).to(device)
    rewards_np = np.array(rewards, dtype=np.float32)
    hx_v = torch.FloatTensor(np.concatenate(hx, axis=1)).to(device)
    cx_v = torch.FloatTensor(np.concatenate(cx, axis=1)).to(device)
    if not_done_idx:
        last_states_v = torch.FloatTensor(np.asarray(last_states)).to(device)
        last_states_hx_v = torch.FloatTensor(np.concatenate(last_states_hx, axis=1)).to(device)
        last_states_cx_v = torch.FloatTensor(np.concatenate(last_states_cx, axis=1)).to(device)
        last_vals_v = net((last_states_v, (last_states_hx_v, last_states_cx_v)))[1]
        last_vals_np = last_vals_v.data.cpu().numpy()[:, 0]
        rewards_np[not_done_idx] += GAMMA ** REWARD_STEPS * last_vals_np

    ref_vals_v = torch.FloatTensor(rewards_np).to(device)   
    return states_v, actions_t, ref_vals_v, hx_v, cx_v


def test_model(env, net, device, episodes=5):
    with torch.no_grad():
        total_reward = 0.0
        for _ in range(episodes):
            noop_action_count = 0
            pre_action = -1
            obs, _ = env.reset()
            hx = torch.zeros((net.lstm.num_layers, 1, net.lstm.hidden_size), dtype=torch.float32).to(device)
            cx = torch.zeros((net.lstm.num_layers, 1, net.lstm.hidden_size), dtype=torch.float32).to(device)
            while True:
                obs_v = ptan.agent.default_states_preprocessor([obs]).to(device)
                logits_v, _, (hx, cx) = net((obs_v, (hx, cx)))
                probs_v = F.softmax(logits_v, dim=1)
                probs = probs_v.data.cpu().numpy()
                action = np.argmax(probs)
                if action == 0 and pre_action == action:  # Noop
                    noop_action_count += 1
                    if noop_action_count > 30:
                        break
                else:
                    noop_action_count = 0
                pre_action = action
                obs, reward, done, trunc, _ = env.step(action)
                total_reward += reward
                if done or trunc:
                    break
    return total_reward / episodes


def optimized_states_preprocessor(states):
    """
    Convert list of states into the form suitable for model.
    :param states: list of numpy arrays with states
    :return: torch.Tensor
    """
    if len(states) == 1:
        np_states = np.expand_dims(states[0], 0)
    else:
        np_states = np.asarray([np.asarray(s) for s in states])
    return torch.from_numpy(np_states)



def select_device(args):
    if args.cuda and torch.cuda.is_available():
        return torch.device("cuda")
    elif torch.backends.mps.is_available() and args.cuda:
        return torch.device("mps")
    return torch.device("cpu")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--cuda", default=True, action="store_true", help="Enable cuda")
    parser.add_argument("-n", "--name", default="fishingderby", required=False, help="Name of the run")
    args = parser.parse_args()
    device = select_device(args)

    save_path = os.path.join("saves", "a2c-lstm-conv" + args.name)
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    envs = [common.wrap_dqn(gym.make("ALE/FishingDerby-v5", obs_type='rgb', frameskip=4, repeat_action_probability=0.0), episodic_life=False) for _ in range(NUM_ENVS)]
    test_env = common.wrap_dqn(gym.make("ALE/FishingDerby-v5", obs_type='rgb', frameskip=4, repeat_action_probability=0.0), episodic_life=False)
    writer = SummaryWriter(comment="-a2c-lstm-conv_" + args.name)

    net = model.AtariA2C(envs[0].observation_space.shape, envs[0].action_space.n).to(device)
    print(net)

    agent = model.PolicyLstmAgent(net, num_envs=NUM_ENVS, apply_softmax=True, device=device, preprocessor=optimized_states_preprocessor)
    exp_source = ptan.experience.ExperienceSourceFirstLastRAW(envs, agent, gamma=GAMMA, steps_count=REWARD_STEPS)
    optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE, eps=1e-3)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=50000, gamma=0.9)

    batch = []
    best_reward = 0
    frame_idx = 1
    start_idx = 0
    train_count = 1

    if os.path.exists(save_path) and len(os.listdir(save_path)) > 0:
        # 增加加载模型的代码
        checkpoints = sorted(filter(lambda x: "epoch" in x, os.listdir(save_path)), key=lambda x: int(x.split('_')[2].split('.')[0]))
        checkpoint = torch.load(os.path.join(save_path, checkpoints[-1]), map_location=device, weights_only=False)
        frame_idx = checkpoint['frame_idx']
        start_idx = checkpoint['start_idx']
        train_count = checkpoint['train_count']
        net.load_state_dict(checkpoint['net'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        scheduler.load_state_dict(checkpoint['scheduler'])
        print("加载模型成功")
        # 打印学习率大小
        print("Learning Rate:", scheduler.get_last_lr()[0])
        print("scheduler last epoch: ", scheduler.last_epoch)

    with common.RewardTracker(writer, stop_reward=10000) as tracker:
        with ptan.common.utils.TBMeanTracker(writer, batch_size=10) as tb_tracker:
            for step_idx, exp in enumerate(exp_source):
                batch.append(exp)

                new_rewards = exp_source.pop_total_rewards()
                if new_rewards:
                    if tracker.reward(new_rewards[0], step_idx + start_idx):
                        break

                if len(batch) < BATCH_SIZE:
                    continue


                states_v, actions_t, vals_ref_v, hx_v, cx_v = unpack_batch(batch, net, device=device)
                batch.clear()

                optimizer.zero_grad()
                logits_v, value_v, _ = net((states_v, (hx_v, cx_v)))
                loss_value_v = F.mse_loss(value_v.squeeze(-1), vals_ref_v)

                log_prob_v = F.log_softmax(logits_v, dim=1)
                adv_v = vals_ref_v - value_v.squeeze(-1).detach()
                log_prob_actions_v = adv_v * log_prob_v[range(BATCH_SIZE), actions_t]
                loss_policy_v = -log_prob_actions_v.mean()

                prob_v = F.softmax(logits_v, dim=1)
                entropy_loss_v = ENTROPY_BETA * (prob_v * log_prob_v).sum(dim=1).mean()

                loss_policy_v.backward(retain_graph=True)
                grads = np.concatenate([p.grad.data.cpu().numpy().flatten()
                                        for p in net.parameters()
                                        if p.grad is not None])

                loss_v = entropy_loss_v + loss_value_v
                loss_v.backward()
                nn_utils.clip_grad_norm_(net.parameters(), CLIP_GRAD)
                optimizer.step()
                loss_v += loss_policy_v
                train_count += 1
                scheduler.step()

                if train_count % 200 == 0:
                    # Test the model
                    net.eval()
                    test_reward = test_model(test_env, net, device=device, episodes=2)
                    net.train()
                    print(f"Test reward: {test_reward:.2f}")
                    common.save_best_model(test_reward, net.state_dict(), save_path, "a2c-lstm-best", keep_best=10)

                if train_count % SAVE_ITERS == 0:
                    checkpoint = {
                        "net": net.state_dict(),
                        "optimizer": optimizer.state_dict(),
                        "frame_idx": frame_idx + start_idx,
                        "start_idx": step_idx + start_idx,
                        "scheduler": scheduler.state_dict(),
                        "train_count": train_count
                    }
                    common.save_checkpoints(frame_idx, checkpoint, save_path, "a2c-lstm", keep_last=3)


                tb_tracker.track("advantage",       adv_v, step_idx + start_idx)
                tb_tracker.track("values",          value_v, step_idx + start_idx)
                tb_tracker.track("batch_rewards",   vals_ref_v, step_idx + start_idx)
                tb_tracker.track("loss_entropy",    entropy_loss_v, step_idx + start_idx)
                tb_tracker.track("loss_policy",     loss_policy_v, step_idx + start_idx)
                tb_tracker.track("loss_value",      loss_value_v, step_idx + start_idx)
                tb_tracker.track("loss_total",      loss_v, step_idx + start_idx)
                tb_tracker.track("grad_l2",         np.sqrt(np.mean(np.square(grads))), step_idx + start_idx)
                tb_tracker.track("grad_max",        np.max(np.abs(grads)), step_idx + start_idx)
                tb_tracker.track("grad_var",        np.var(grads), step_idx + start_idx)